from pandas import concat
from model.task import filter_numeric_col
from utils.format_util import cast_float
import numpy as np

def run(df, taskId, columns, n_in=0, n_out=0, drop=True):
    results = {}
    numeric_cols, non_numeric_cols = filter_numeric_col(taskId)

    data, names = [], []
    # i: n_in, n_in-1, ..., 1
    # 代表t-n_in, ... ,t-1
    for i in range(n_in, 0, -1):
        data.append(df[columns].shift(i))
        names += [('%s(t-%d)' % (column, i)) for column in columns]
    for i in range(0, n_out+1):
        data.append(df[columns].shift(-i))
        if i == 0:
            names += [('%s(t)' % (column)) for column in columns]
        else:
            names += [('%s(t+%d)' % (column, i)) for column in columns]

    other_columns = [i for i in df.columns.tolist() if i not in columns]
    for column in other_columns:
        data.append(df[column])
    output = concat(data, axis=1)
    columns = names + other_columns
    output.columns = columns
    output_drop = output.dropna()
    numeric_cols = []
    for col in columns:
        if col not in non_numeric_cols and col != "_record_id_":
            numeric_cols.append(col)

    corr = output_drop[numeric_cols].corr().values.reshape(-1).tolist()
    if not np.isnan(corr[0]):
        results["corr"] = cast_float(corr)
        results["cols"] = numeric_cols
    if drop:
        return output_drop, results
    else:
        return output, results